package com.miqt.plugin.inject;

import com.android.build.api.transform.DirectoryInput;
import com.android.build.api.transform.JarInput;
import com.android.build.api.transform.TransformInput;
import com.android.build.api.transform.TransformInvocation;
import com.android.build.gradle.BaseExtension;
import com.miqt.asm.method_hook.BasePlugin;
import com.miqt.plugin.inject.compile.ClassRewriter;
import com.miqt.plugin.inject.compile.FileLog;
import com.miqt.plugin.inject.hook.HookClassesConfig;

import org.gradle.api.NamedDomainObjectContainer;
import org.gradle.api.Project;
import org.gradle.api.plugins.ExtensionAware;
import org.jetbrains.annotations.NotNull;

import java.io.File;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.jar.JarEntry;

public class InjectPlugin extends BasePlugin<AutotrackExtension> {

    private ClassRewriter mClassRewriter;
    private Collection<URL> mAndroidJars;

    @Override
    public void apply(@NotNull Project project) {
        super.apply(project);
        project.afterEvaluate(project1 -> {
            getAndroidJar(project);
        });
    }

    private void getAndroidJar(Project project) {
        try {
            BaseExtension android = (BaseExtension) project.getExtensions().findByName("android");
            if (android == null) {
                return;
            }
            List<File> files = android.getBootClasspath();
            if (files.isEmpty()) {
                throw new RuntimeException("Inject Plugin: get android.jar failed");
            }
            List<URL> androidJars = new ArrayList<>();
            for (File file : files) {
                androidJars.add(file.toURL());
            }
            mAndroidJars = androidJars;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("Inject Plugin: get android.jar failed");
        }
    }

    @Override
    public String getName() {
        return "inject_plugin";
    }

    @Override
    protected void beginTransform(TransformInvocation transformInvocation) {
        super.beginTransform(transformInvocation);
        try {
            ArrayList<URL> urlList = new ArrayList<>();
            for (TransformInput input : transformInvocation.getInputs()) {
                for (DirectoryInput directoryInput : input.getDirectoryInputs()) {
                    urlList.add(directoryInput.getFile().toURL());
                }

                for (JarInput jarInput : input.getJarInputs()) {
                    urlList.add(jarInput.getFile().toURL());
                }
            }
            urlList.addAll(mAndroidJars);
            URL[] urlArray = new URL[urlList.size()];
            urlList.toArray(urlArray);
            URLClassLoader classLoader = new URLClassLoader(urlArray);
            HookClassesConfig.init(classLoader.loadClass("com.miqt.plugin.inject.HookTargets"));
            mClassRewriter = new ClassRewriter(new FileLog(getLogger()), classLoader, getExtension().getExcludePackages(), getExtension().isExcludeOfficialPackages());
        } catch (Throwable e) {
            getLogger().log(e);
        }

    }

    @Override
    public AutotrackExtension initExtension() {
        return new AutotrackExtension();
    }

    @Override
    public byte[] transform(byte[] classBytes, File classFile) {
        byte[] res = mClassRewriter.visitClassBytes(classBytes);
        if (res == null) {
            return classBytes;
        } else {
            return res;
        }
    }

    @Override
    public byte[] transformJar(byte[] classBytes, File jarFile, JarEntry entry) {
        byte[] res = mClassRewriter.visitClassBytes(classBytes);
        if (res == null) {
            return classBytes;
        } else {
            return res;
        }
    }
}
